from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
import pymysql
from datetime import datetime, timedelta
import requests
import json
from config import Config
import hashlib
from typing import Dict, Any, Optional
from functools import wraps

app = Flask(__name__)
CORS(app)
app.config.from_object(Config)
jwt = JWTManager(app)

def get_db():
    return pymysql.connect(
        host=Config.MYSQL_HOST,
        port=Config.MYSQL_PORT,
        user=Config.MYSQL_USER,
        password=Config.MYSQL_PASSWORD,
        db=Config.MYSQL_DB,
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

def handle_db_error(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except Exception as e:
            return jsonify({'error': str(e)}), 400
    return decorated_function

# 用户认证相关API
@app.route('/api/register', methods=['POST'])
@handle_db_error
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    wechat_user_id = data.get('wechat_user_id')
    
    if not username or not password:
        return jsonify({'error': '用户名和密码不能为空'}), 400
    
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                'INSERT INTO users (username, password, wechat_user_id) VALUES (%s, %s, %s)',
                (username, hashed_password, wechat_user_id)
            )
            conn.commit()
    return jsonify({'message': '注册成功'}), 201

@app.route('/api/login', methods=['POST'])
@handle_db_error
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                'SELECT id, username FROM users WHERE username = %s AND password = %s',
                (username, hashed_password)
            )
            user = cursor.fetchone()
            
            if user:
                access_token = create_access_token(identity=user['id'])
                return jsonify({
                    'access_token': access_token,
                    'user': user
                })
            return jsonify({'error': '用户名或密码错误'}), 401

# 基金相关API
@app.route('/api/fund/<code>', methods=['GET'])
@handle_db_error
def get_fund_detail(code: str):
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('SELECT * FROM funds WHERE code = %s', (code,))
            fund = cursor.fetchone()
            
            if not fund:
                # 如果基金不存在，从API获取并保存
                response = requests.get(Config.FUND_API_URL.format(code))
                if response.status_code == 200:
                    data = response.text
                    # 解析基金数据
                    name = data.split('fS_name = "')[1].split('"')[0]
                    cursor.execute(
                        'INSERT INTO funds (code, name) VALUES (%s, %s)',
                        (code, name)
                    )
                    conn.commit()
                    fund = {'code': code, 'name': name}
                else:
                    return jsonify({'error': '基金不存在'}), 404
            
            return jsonify(fund)

@app.route('/api/fund/<code>/net_values', methods=['GET'])
@handle_db_error
def get_fund_net_values(code: str):
    start_date = request.args.get('start_date')
    end_date = request.args.get('end_date')
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('SELECT id FROM funds WHERE code = %s', (code,))
            fund = cursor.fetchone()
            if not fund:
                return jsonify({'error': '基金不存在'}), 404
            
            query = 'SELECT * FROM fund_net_values WHERE fund_id = %s'
            params = [fund['id']]
            
            if start_date:
                query += ' AND date >= %s'
                params.append(start_date)
            if end_date:
                query += ' AND date <= %s'
                params.append(end_date)
            
            query += ' ORDER BY date DESC'
            
            cursor.execute(query, params)
            net_values = cursor.fetchall()
            return jsonify(net_values)

# 用户自选基金相关API
@app.route('/api/user/funds', methods=['GET'])
@jwt_required()
@handle_db_error
def get_user_funds():
    user_id = get_jwt_identity()
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('''
                SELECT uf.*, f.code, f.name 
                FROM user_funds uf 
                JOIN funds f ON uf.fund_id = f.id 
                WHERE uf.user_id = %s
            ''', (user_id,))
            funds = cursor.fetchall()
            return jsonify(funds)

@app.route('/api/user/fund', methods=['POST'])
@jwt_required()
@handle_db_error
def add_user_fund():
    user_id = get_jwt_identity()
    data = request.get_json()
    fund_code = data.get('fund_code')
    buy_date = data.get('buy_date')
    expected_rate = data.get('expected_rate')
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            # 获取基金ID
            cursor.execute('SELECT id FROM funds WHERE code = %s', (fund_code,))
            fund = cursor.fetchone()
            if not fund:
                return jsonify({'error': '基金不存在'}), 404
            
            # 添加自选
            cursor.execute('''
                INSERT INTO user_funds (user_id, fund_id, buy_date, expected_rate)
                VALUES (%s, %s, %s, %s)
            ''', (user_id, fund['id'], buy_date, expected_rate))
            conn.commit()
            return jsonify({'message': '添加成功'})

# 监控规则相关API
@app.route('/api/monitor/rules', methods=['GET'])
@jwt_required()
@handle_db_error
def get_monitor_rules():
    user_id = get_jwt_identity()
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('''
                SELECT mr.*, f.code, f.name 
                FROM monitor_rules mr 
                JOIN funds f ON mr.fund_id = f.id 
                WHERE mr.user_id = %s
            ''', (user_id,))
            rules = cursor.fetchall()
            return jsonify(rules)

@app.route('/api/monitor/rule', methods=['POST'])
@jwt_required()
@handle_db_error
def add_monitor_rule():
    user_id = get_jwt_identity()
    data = request.get_json()
    
    with get_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('''
                INSERT INTO monitor_rules 
                (user_id, fund_id, rule_type, threshold, start_date, end_date)
                VALUES (%s, %s, %s, %s, %s, %s)
            ''', (
                user_id,
                data['fund_id'],
                data['rule_type'],
                data['threshold'],
                data.get('start_date'),
                data.get('end_date')
            ))
            conn.commit()
            return jsonify({'message': '添加成功'})

if __name__ == '__main__':
    app.run(debug=True) 